"""
Streamlit RAG 应用
"""
import os
import time
import streamlit as st
from dotenv import load_dotenv
from config import DOCUMENTS_DIR
from vector_store import VectorStoreManager
from document_processor import DocumentProcessor

# 设置页面标题
st.set_page_config(
    page_title="智能问答系统",
    page_icon="🤖",
    layout="wide"
)

# 添加标题
st.title("📚 智能问答系统")

def initialize_session_state():
    """初始化会话状态"""
    if 'vector_store' not in st.session_state:
        # 加载环境变量
        load_dotenv()
        
        # 检查 API key
        if not os.getenv("OPENAI_API_KEY"):
            st.error("错误: 未设置 OPENAI_API_KEY 环境变量")
            st.stop()
            
        # 初始化向量存储
        st.session_state.vector_store = VectorStoreManager()

        # 初始化文档处理器
        st.session_state.doc_processor = DocumentProcessor()
        
    # 初始化对话历史
    if 'messages' not in st.session_state:
        st.session_state.messages = []
        
    # 初始化索引
    if 'index' not in st.session_state:
        if os.path.exists(DOCUMENTS_DIR):
            with st.spinner("正在创建文档索引..."):
                st.session_state.index = st.session_state.vector_store.create_index()

def add_message(role, content):
    """添加消息到对话历史"""
    st.session_state.messages.append({"role": role, "content": content})

def display_chat_history():
    """显示对话历史"""
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

def process_uploaded_files(files):
    """处理上传的文件"""
    if files:
        with st.spinner("正在处理上传的文件..."):
            # 创建新的索引
            new_index = st.session_state.vector_store.create_index(files=files)
            
            # 如果已经存在索引，则合并
            if 'index' in st.session_state:
                # TODO: 实现索引合并逻辑
                st.session_state.index = new_index
            else:
                st.session_state.index = new_index
                
        st.success(f"成功上传 {len(files)} 个文件！")
        
def main():
    """主程序"""
    initialize_session_state()

    # 创建侧边栏
    with st.sidebar:
        st.header("💡 使用说明")
        st.markdown("""
        1. 系统已自动加载文档目录下的所有文件
        2. 您可以通过下方的文件上传功能添加新文档
        3. 在对话框中输入您的问题
        """)
        
        # 文档统计
        st.subheader("文档统计")
        doc_stats = st.session_state.doc_processor.get_document_stats()
        st.write(f"总文件数: {doc_stats['total_files']}")
        st.write(f"有效文件数: {doc_stats['valid_files']}")
        
        # 显示文件类型统计
        if doc_stats['file_types']:
            st.write("文件类型分布:")
            for ext, count in doc_stats['file_types'].items():
                st.write(f"- {ext}: {count}个文件")

        # 总嵌入数
        st.subheader("总嵌入数")
        st.write(f"总嵌入数: {st.session_state.vector_store.get_collection_stats()}")

        # 添加文件上传功能
        st.header("📤 上传文件")
        uploaded_files = st.file_uploader(
            "选择要上传的文件",
            accept_multiple_files=True,
            type=["txt", "md", "pdf", "html"],
        )
        
        if uploaded_files is not None:
            process_uploaded_files(uploaded_files)
            
    # 显示对话历史
    display_chat_history()
    
    # 创建问答界面
    if prompt := st.chat_input("请输入您的问题..."):
        # 检查是否已创建索引
        if 'index' not in st.session_state:
            st.warning("请先上传文件或确保文档目录中有文件。")
            st.stop()
            
        # 添加用户问题到对话历史
        add_message("user", prompt)
        
        # 显示最新的用户问题
        with st.chat_message("user"):
            st.markdown(prompt)
            
        # 显示助手回复
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            
            # 获取流式响应
            try:
                response = st.session_state.vector_store.query(
                    st.session_state.index,
                    prompt
                )
                
                # 逐字显示回复
                for chunk in response.response_gen:
                    full_response += chunk
                    message_placeholder.markdown(full_response + "▌")
                    time.sleep(0.01)
                    
                # 显示完整回复
                message_placeholder.markdown(full_response)
                
                # 添加助手回复到对话历史
                add_message("assistant", full_response)
                
            except Exception as e:
                st.error(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main() 